// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/least_greatest.h>

#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/functions/lib/simd_comparison_util.h>
#include <pollux/functions/sparksql/comparisons.h>
#include <pollux/type/type.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename Cmp, TypeKind kind>
        class ComparisonFunction final : public exec::VectorFunction {
            using T = typename TypeTraits<kind>::NativeType;

            bool supportsFlatNoNullsFastPath() const override {
                return true;
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                const Cmp cmp;
                context.ensure_writable(rows, BOOLEAN(), result);
                result->clear_nulls(rows);
                if constexpr (
                    kind == TypeKind::TINYINT || kind == TypeKind::SMALLINT ||
                    kind == TypeKind::INTEGER || kind == TypeKind::BIGINT) {
                    if ((args[0]->is_flat_encoding() || args[0]->is_constant_encoding()) &&
                        (args[1]->is_flat_encoding() || args[1]->is_constant_encoding()) &&
                        rows.isAllSelected()) {
                        applySimdComparison<T, Cmp>(rows, args, result);
                        return;
                    }
                }
                if (shouldApplyAutoSimdComparison<T, T>(rows, args)) {
                    applyAutoSimdComparison<T, T, Cmp>(rows, args, context, result);
                    return;
                }

                auto *flatResult = result->as_unchecked<FlatVector<bool> >();

                if (args[0]->is_flat_encoding() && args[1]->is_flat_encoding()) {
                    // Fast path for (flat, flat).
                    const auto *rawA = args[0]->as_unchecked<FlatVector<T> >()->rawValues();
                    const auto *rawB = args[1]->as_unchecked<FlatVector<T> >()->rawValues();
                    rows.applyToSelected(
                        [&](vector_size_t i) { flatResult->set(i, cmp(rawA[i], rawB[i])); });
                } else if (args[0]->is_constant_encoding() && args[1]->is_flat_encoding()) {
                    // Fast path for (const, flat).
                    auto constant = args[0]->as_unchecked<ConstantVector<T> >()->value_at(0);
                    const auto *rawValues =
                            args[1]->as_unchecked<FlatVector<T> >()->rawValues();
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(i, cmp(constant, rawValues[i]));
                    });
                } else if (args[0]->is_flat_encoding() && args[1]->is_constant_encoding()) {
                    // Fast path for (flat, const).
                    const auto *rawValues =
                            args[0]->as_unchecked<FlatVector<T> >()->rawValues();
                    auto constant = args[1]->as_unchecked<ConstantVector<T> >()->value_at(0);
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(i, cmp(rawValues[i], constant));
                    });
                } else {
                    // Path if one or more arguments are encoded.
                    exec::DecodedArgs decodedArgs(rows, args, context);
                    auto decodedA = decodedArgs.at(0);
                    auto decodedB = decodedArgs.at(1);
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(
                            i, cmp(decodedA->value_at<T>(i), decodedB->value_at<T>(i)));
                    });
                }
            }
        };

        // ComparisonFunction instance for bool as it uses compact representation
        template<typename Cmp>
        class BoolComparisonFunction final : public exec::VectorFunction {
            bool supportsFlatNoNullsFastPath() const override {
                return true;
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                context.ensure_writable(rows, BOOLEAN(), result);
                auto *flatResult = result->as_unchecked<FlatVector<bool> >();
                const Cmp cmp;

                if (args[0]->is_flat_encoding() && args[1]->is_flat_encoding()) {
                    // Fast path for (flat, flat).
                    auto rawA = args[0]
                            ->as_unchecked<FlatVector<bool> >()
                            ->template mutableRawValues<uint64_t>();
                    auto rawB = args[1]
                            ->as_unchecked<FlatVector<bool> >()
                            ->template mutableRawValues<uint64_t>();
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(
                            i, cmp(bits::isBitSet(rawA, i), bits::isBitSet(rawB, i)));
                    });
                } else if (args[0]->is_constant_encoding() && args[1]->is_flat_encoding()) {
                    // Fast path for (const, flat).
                    auto constant = args[0]->as_unchecked<ConstantVector<bool> >()->value_at(0);
                    auto rawValues = args[1]
                            ->as_unchecked<FlatVector<bool> >()
                            ->template mutableRawValues<uint64_t>();
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(i, cmp(constant, bits::isBitSet(rawValues, i)));
                    });
                } else if (args[0]->is_flat_encoding() && args[1]->is_constant_encoding()) {
                    // Fast path for (flat, const).
                    auto rawValues = args[0]
                            ->as_unchecked<FlatVector<bool> >()
                            ->template mutableRawValues<uint64_t>();
                    auto constant = args[1]->as_unchecked<ConstantVector<bool> >()->value_at(0);
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(i, cmp(bits::isBitSet(rawValues, i), constant));
                    });
                } else {
                    // Fast path if one or more arguments are encoded.
                    exec::DecodedArgs decodedArgs(rows, args, context);
                    auto decodedA = decodedArgs.at(0);
                    auto decodedB = decodedArgs.at(1);
                    rows.applyToSelected([&](vector_size_t i) {
                        flatResult->set(
                            i, cmp(decodedA->value_at<bool>(i), decodedB->value_at<bool>(i)));
                    });
                }
            }
        };

        template<template <typename> class Cmp, typename StdCmp>
        std::shared_ptr<exec::VectorFunction> makeImpl(
            const std::string &functionName,
            const std::vector<exec::VectorFunctionArg> &args) {
            POLLUX_CHECK_EQ(args.size(), 2);
            for (size_t i = 1; i < args.size(); i++) {
                POLLUX_CHECK(*args[i].type == *args[0].type);
            }
            switch (args[0].type->kind()) {
#define SCALAR_CASE(kind)                            \
  case TypeKind::kind:                               \
    return std::make_shared<ComparisonFunction<      \
        Cmp<TypeTraits<TypeKind::kind>::NativeType>, \
        TypeKind::kind>>();
                SCALAR_CASE(REAL)
                SCALAR_CASE(DOUBLE)
                SCALAR_CASE(HUGEINT)
                SCALAR_CASE(VARCHAR)
                SCALAR_CASE(VARBINARY)
                SCALAR_CASE(TIMESTAMP)
#undef SCALAR_CASE
#define STD_SCALAR_CASE(kind) \
  case TypeKind::kind:        \
    return std::make_shared<ComparisonFunction<StdCmp, TypeKind::kind>>();
                STD_SCALAR_CASE(TINYINT)
                STD_SCALAR_CASE(SMALLINT)
                STD_SCALAR_CASE(INTEGER)
                STD_SCALAR_CASE(BIGINT)
#undef STDSCALAR_CASE
                case TypeKind::BOOLEAN:
                    return std::make_shared<BoolComparisonFunction<
                        Cmp<TypeTraits<TypeKind::BOOLEAN>::NativeType> > >();
                default:
                    POLLUX_NYI(
                        "{} does not support arguments of type {}",
                        functionName,
                        args[0].type->kind());
            }
        }

        template<TypeKind kind>
        void applyTyped(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            DecodedVector *decodedLhs,
            DecodedVector *decodedRhs,
            exec::EvalCtx &context,
            FlatVector<bool> *flatResult) {
            using T = typename TypeTraits<kind>::NativeType;

            Equal<T> equal;
            if (!args[0]->may_have_nulls() && !args[1]->may_have_nulls()) {
                // When there is no nulls, it reduces to normal equality function
                rows.applyToSelected([&](vector_size_t i) {
                    flatResult->set(
                        i, equal(decodedLhs->value_at<T>(i), decodedRhs->value_at<T>(i)));
                });
            } else {
                // (isnull(a) AND isnull(b)) || (a == b)
                // When DecodedVector::nulls() is null it means there are no nulls.
                auto *rawNulls0 = decodedLhs->nulls(&rows);
                auto *rawNulls1 = decodedRhs->nulls(&rows);
                rows.applyToSelected([&](vector_size_t i) {
                    auto isNull0 = rawNulls0 && bits::isBitNull(rawNulls0, i);
                    auto isNull1 = rawNulls1 && bits::isBitNull(rawNulls1, i);
                    flatResult->set(
                        i,
                        (isNull0 || isNull1)
                            ? isNull0 && isNull1
                            : equal(decodedLhs->value_at<T>(i), decodedRhs->value_at<T>(i)));
                });
            }
        }

        class EqualtoNullSafe final : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                exec::DecodedArgs decodedArgs(rows, args, context);

                DecodedVector *decodedLhs = decodedArgs.at(0);
                DecodedVector *decodedRhs = decodedArgs.at(1);
                context.ensure_writable(rows, BOOLEAN(), result);
                auto *flatResult = result->as_unchecked<FlatVector<bool> >();
                flatResult->mutableRawValues<int64_t>();

                POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                    applyTyped,
                    args[0]->type_kind(),
                    rows,
                    args,
                    decodedLhs,
                    decodedRhs,
                    context,
                    flatResult);
            }
        };
    } // namespace

    std::shared_ptr<exec::VectorFunction> makeEqualTo(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<Equal, std::equal_to<> >(functionName, args);
    }

    std::shared_ptr<exec::VectorFunction> makeLessThan(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<Less, std::less<> >(functionName, args);
    }

    std::shared_ptr<exec::VectorFunction> makeGreaterThan(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<Greater, std::greater<> >(functionName, args);
    }

    std::shared_ptr<exec::VectorFunction> makeLessThanOrEqual(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<LessOrEqual, std::less_equal<> >(functionName, args);
    }

    std::shared_ptr<exec::VectorFunction> makeGreaterThanOrEqual(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<GreaterOrEqual, std::greater_equal<> >(functionName, args);
    }

    std::shared_ptr<exec::VectorFunction> makeEqualToNullSafe(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig & /*config*/) {
        static const auto kEqualtoNullSafeFunction =
                std::make_shared<EqualtoNullSafe>();
        return kEqualtoNullSafeFunction;
    }
} // namespace kumo::pollux::functions::sparksql
